
import dashscope
from typing import List, Sequence, Optional, Any

from langchain_core.callbacks import Callbacks
from langchain_core.documents import Document
import requests
from langchain.retrievers.document_compressors.base import (
    BaseDocumentCompressor,
)

from xinference_client.client.restful.restful_client import Client
from langchain.chains import FlareChain

from config.LLMConfig import XINFERENCE_ENDPOINT, XINFERENCE_RERANK_MODEL


class DashScopeRerankCompressor(BaseDocumentCompressor):
    model: str = "gte-rerank-v2"
    top_n:  int = 3

    def __init__(self, model: str = "gte-rerank-v2", top_n: int = 3):
        """
        初始化 DashScopeTextReRankCompressor。

        参数:
            api_key (str): DashScope API 密钥
            model (str): TextReRank 模型名称，默认为 "text-rerank-model"
        """
        # 调用父类 BaseDocumentCompressor 的 __init__ 方法
        super().__init__()
        self.model = model
        self.top_n = top_n
    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        # 提取文档内容
        docs_text = [doc.page_content for doc in documents]

        # 调用DashScope Rerank API
        response = dashscope.TextReRank.call(
            model=self.model,
            query=query,
            documents=docs_text,
            top_n=self.top_n
        )

        # 按得分排序文档
        sorted_docs = sorted(
            zip(documents, response.output.results),
            key=lambda x: x[1].relevance_score,
            reverse=True
        )
        return [doc for doc, _ in sorted_docs[:self.top_n]]


class XinferenceRerankCompressor(BaseDocumentCompressor):
    model: str = "gte-rerank-v2"
    top_n:  int = 3
    endpoint: str = ""
    def __init__(self,endpoint:str, model: str = "gte-rerank-v2", top_n: int = 3):
        """
        初始化 DashScopeTextReRankCompressor。

        参数:
            api_key (str): DashScope API 密钥
            model (str): TextReRank 模型名称，默认为 "text-rerank-model"
        """
        # 调用父类 BaseDocumentCompressor 的 __init__ 方法
        super().__init__()
        self.model = model
        self.top_n = top_n
        self.endpoint = endpoint

    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        docs_text = [doc.page_content for doc in documents]
        client = Client(XINFERENCE_ENDPOINT)
        model = client.get_model(XINFERENCE_RERANK_MODEL)

        response = model.rerank(docs_text, query)
        print(response)
        scores = [result["relevance_score"] for result in response["results"]]
        sorted_docs = sorted(
            zip(documents, scores),
            key=lambda x: x[1],
            reverse=True
        )
        return [doc for doc, _ in sorted_docs[:self.top_n]]


